import logging

from composite_task_queue import CompositeTaskQueue
from robot_controller import RobotController
from robots_manager import RobotsManager
from dispatcher import Dispatcher
from router import Router
from scheduler import Scheduler


class TransportationControlSystem(object):
    EVENT_NEW_TASK_CREATED = 0
    EVENT_ROBOT_POS_UPDATED = 1
    
    def __init__(self, world):
        self.initLogger()
        self._world = world
        
        self._ctask_queue = CompositeTaskQueue()
        self._buffered_ctasks = []
        self._robots_manager = RobotsManager(world)
        
        self._timer = world.getTimer()
        self._map = world.getMap()
        self._router = Router(self._map)
        self._dispatcher = Dispatcher(self._timer, self._map, self._router)
        self._scheduler = Scheduler(self._map, self._robots_manager, self._timer)
        self._robot2index = {}
        
    def initLogger(self):
        # create logger
        self._logger = logging.getLogger('MyTCS.TCS')
        
    def getRouter(self):
        return self._router
        
    def getDispatcher(self):
        return self._dispatcher
        
    def getScheduler(self):
        return self._scheduler            

    def addCompositeTask(self, ctask):
        """
        向复合任务队列中添加一个新的复合任务。
        :param ctask: 欲添加的复合任务。
        :param delay: 自本函数调用后多长时间添加，默认为0，即立刻添加。实际上，本函数并不直接将ctask添加进复合任务队列，而是缓存起来，
                      到delay所指示的时间后，再从缓存中添加到复合任务队列。见updateCompositeTaskQueueFromBuffer函数。
        """
        lo = 0
        hi = len(self._buffered_ctasks)
        delay = ctask.getCreateTime() - self._timer.getTime()
        if delay < 0:
            raise ValueError('The create time should be larger than or equal to the current time!')
        while lo < hi:
            mid = (lo+hi)//2
            mid_delay = self._buffered_ctasks[mid][0]
            if mid_delay > delay: 
                hi = mid
            else:
                lo = mid + 1
        self._buffered_ctasks.insert(lo, [delay, ctask])
        
    def isTaskQueueEmpty(self):
        """
        检查复合任务队列中是否为空。
        :return: True-空，False-非空。
        """
        return (not self._buffered_ctasks) and self._ctask_queue.isEmpty()
           
    def allTasksCompleted(self, creator='user'):
        for buf in self._buffered_ctasks:
            if buf[1].getCreator() == creator:
                next_task = buf[1].getNextUncompletedTask()
                last_task = buf[1].getLastUncompletedTask()
                if next_task is not None:
                    self._logger.debug('buffered user task ({0}, {1})'.format(next_task.getSrc(), last_task.getDst()))
                return False
        for ctask in self._ctask_queue:
            if ctask.getCreator() == creator:
                next_task = ctask.getNextUncompletedTask()
                last_task = ctask.getLastUncompletedTask()
                if next_task is not None:
                    self._logger.debug('user task ({0}, {1})'.format(next_task.getSrc(), last_task.getDst()))
                return False
        return True
        
    def scanCompositeTaskQueue(self):
        """
        在任务执行过程中，某些CompositeTask中包含的Task数量会随着任务完成而减为0.此函数将这样的CompositeTask从复合任务队列中移除。
        此外，某些CompositeTask中的子任务将会完成，设置相应的状态。
        """
        i = 0
        ctask_queue_updated = False
        while i < self._ctask_queue.size():
            #print(i)
            ctask = self._ctask_queue[i]
            if ctask.isEmpty():
                self._ctask_queue.removeAt(i)
                ctask_queue_updated = True
                #print('remove task ({0}, {1}) A'.format(task.getSrc(), task.getDst()))
                i -= 1
            elif ctask.isAssigned():
                controller = ctask.getAssignedRobotController()
                last_waypoint = controller.getLastWaypoint()
                task = ctask.getNextUncompletedTask()
                #self._logger.debug('robot {0}, last waypoint {1}, task ({2}, {3})'.format(controller.getControlledRobotID(), last_waypoint, task.getSrc(), task.getDst()))
                if last_waypoint == task.getDst():
                    # FixMe: we might need to perform some actions here
                    #print('remove task ({0}, {1}) B'.format(task.getSrc(), task.getDst()))
                    controller.setCurrentSimpleTaskCompleted()
                    del self._robot2index[controller.getControlledRobotID()]
                    if ctask.isEmpty():
                        self._ctask_queue.removeAt(i)
                        ctask_queue_updated = True
                        i -= 1
            i += 1
        if ctask_queue_updated:
            self._updateRobot2Index()
                
    def updateCompositeTaskQueueFromBuffer(self, dt):
        """
        根据缓存的复合任务更新复合任务队列。
        :param dt: 步进时间。缓存以dt时间为步长被周期性的扫描。一旦到达缓存中某条任务指定的添加时间
                   （由addCompositeTask函数的delay参数指定），即将其添加至复合任务队列。
        """
        i = 0
        while i < len(self._buffered_ctasks):
            self._buffered_ctasks[i][0] -= dt
            if self._buffered_ctasks[i][0] <= 0:
                self._ctask_queue.append(self._buffered_ctasks[i][1], sort=True)
                self._logger.debug('Time={0:.2f} TCS update task queue from buffer, new task ({1}, {2}) arrived.'.\
                                   format(self._timer.getTime(), self._buffered_ctasks[i][1][0].getSrc(), self._buffered_ctasks[i][1][0].getDst()))
                print('Time={0:.2f} TCS update task queue from buffer, new task ({1}, {2}) arrived.'.\
                                   format(self._timer.getTime(), self._buffered_ctasks[i][1][0].getSrc(), self._buffered_ctasks[i][1][0].getDst()))
                del self._buffered_ctasks[i]
            else:
                i += 1
                
    def updatePriority(self, task=None):
        """
        更新任务的优先级。
        参考文献：
        Smolic-Rocak N, Bogdan S, Kovacic Z, et al. Time Windows Based Dynamic Routing in Multi-AGV Systems[J]. 
        IEEE Transactions on Automation Science & Engineering, 2009, 7(1):151-155.
        该文献中任务的优先级高低与priority的数值大小是反的，即priority越大，优先级越低。本实现中将其调整为一致了。
        """
        t = self._timer.getTime()

        if task is not None:
            if task.isProcessing() or task.isPending():
                if task.getReleaseTime() == float('inf'):
                    p = float('inf')
                else:
                    p0 = task.getInitialPriority()
                    p = max(2*p0-(task.getDueTime()-t)/(task.getReleaseTime()-t), p0)
                task.setPriority(p)
        else:
            for ctask in self._ctask_queue:
                #print('Trying to update ', ctask.getAssignedRobotController().getControlledRobotID(), '\'s task priority')
                task = ctask.getNextUncompletedTask()
                if task.isProcessing() or task.isPending():
                    if task.getReleaseTime() == float('inf'):
                        p = float('inf')
                    else:
                        p0 = task.getInitialPriority()
                        p = max(2*p0-(task.getDueTime()-t)/(task.getReleaseTime()-t), p0)
                    task.setPriority(p)
        #self._logger.debug('Before sorting ...')
        #self._printPriorities()
        self._ctask_queue.quicksort()
        #self._logger.debug('After sorting ...')
        #self._printPriorities()
    
    def _printPriorities(self):
        ctasks = list(self._ctask_queue)
        for ctask in self._ctask_queue:
            if not ctask.isAssigned():
                continue
            robot_id = ctask.getAssignedRobotController().getControlledRobotID()
            priority = ctask.getNextUncompletedTask().getPriority()
            self._logger.debug('Task assigned to robot {0}, with priority {1}'.format(robot_id, priority))

    def determinePriority(self, task, cur_pos=-1):
        k = self._indexOfProcessingOrPendingTask(-1, reverse=True)
        if k < 0:
            task.setPriority(0)
            return
        sorted_ctasks = list(self._ctask_queue)
        # 新创建的任务总是添加在链表最末尾
        # 此处，首先将新任务移动到最前面
        self._logger.debug('moving task from {0} to {1}'.format(cur_pos, k))
        self._ctask_queue.move(cur_pos, k)
        cur_pos = k
        min_avg_response_time = float('inf')
        min_num_pending_tasks = self._robots_manager.getRobotControllersNum()
        opt_pos = k
        last_task_found = False
        
        while True:
            task.setUnscheduled()
            self._scheduler.backup()
            #self.reschedule(TransportationControlSystem.EVENT_NEW_TASK_CREATED, verbose=False, planning=True)
            pending_ctasks = set(self.pendingCompositeTasks(start=cur_pos+1))
            schedule_results = self._rescheduleTasksWithLowerPriorities(self._robot2index, start=cur_pos, verbose=False)
            pending_ctasks1 = set(self.pendingCompositeTasks(start=cur_pos+1))
            if len(pending_ctasks1 - pending_ctasks) > 0:
                break
                
            avg_response_time, num_pending_tasks = self.avgResponseTime()
            self._logger.debug('Current position {0}, avg response time {1}, # pending tasks {2}'.format(cur_pos, avg_response_time, num_pending_tasks))
            if num_pending_tasks < min_num_pending_tasks:
                min_num_pending_tasks = num_pending_tasks
                min_avg_response_time = avg_response_time
                opt_pos = cur_pos
                self._logger.debug('Optimal position changed to {0}'.format(opt_pos))
            elif (num_pending_tasks == min_num_pending_tasks and \
                  avg_response_time < min_avg_response_time):
                min_avg_response_time = avg_response_time
                opt_pos = cur_pos
                self._logger.debug('Optimal position changed to {0}'.format(opt_pos))
            self._scheduler.restore()

            if last_task_found:
                break
                
            k1 = self._indexOfProcessingOrPendingTask(k-1, reverse=True)
            if k1 < 0:
                self._ctask_queue.move(cur_pos, k)
                self._logger.debug('moving task from {0} to {1}'.format(cur_pos, k))
                cur_pos = k
                last_task_found = True
            else:
                self._ctask_queue.move(cur_pos, k1-1)
                self._logger.debug('moving task from {0} to {1}'.format(cur_pos, k1-1))
                cur_pos = k1-1
                k = k1
        self._ctask_queue.move(cur_pos, opt_pos)
        task.setUnscheduled()
    
    def process(self, dt):
        """
        传输控制系统（TransportationControlSystem）的集中处理方法。
        本方法应被周期性轮询，dt为轮询周期。
        """
        self._robots_manager.process(dt)
        
        self.scanCompositeTaskQueue()
        self.updateCompositeTaskQueueFromBuffer(dt)

        assign_list = self._dispatcher.dispatch(self._ctask_queue, self._robots_manager)

        c = self._robots_manager.getNewTaskArrivedRobot()
        if c is not None:
            self.updatePriority()
            self.reschedule(TransportationControlSystem.EVENT_NEW_TASK_CREATED, verbose=True)
        c = self._robots_manager.getPosUpdatedRobot()
        if c is not None:
            self.reschedule(TransportationControlSystem.EVENT_ROBOT_POS_UPDATED, verbose=True)
            
    def _rescheduleTasks(self, robot2index, start=0, reschedule_lower_priority_tasks=True, verbose=False):
        """
        重新调度sorted_ctasks[start:]中的所有任务。若其中第一个任务，即sorted_ctasks[start]未能成功调度，
        便不再继续调度其后的任务，此时函数返回False。否则返回True。
        """
        sorted_ctasks = list(self._ctask_queue)
        schedule_results = {}
        if reschedule_lower_priority_tasks:
            ctask_idx = self._indexOfProcessingTask(start + 1)
            to_schedule = [start]

            while ctask_idx < len(sorted_ctasks) and ctask_idx >= 0:
                to_schedule += [ctask_idx]
                ctask_idx = self._indexOfProcessingTask(ctask_idx + 1)
        else:
            to_schedule = [start]
        self._logger.debug('to_schedule = {}'.format(to_schedule))
        
        start_robot_id = None
        self._scheduler.backup()
        for ii in to_schedule:
            ctask = sorted_ctasks[ii]
            if ctask.isEmpty() or (not ctask.isAssigned()):
                if ii == start:
                    return {}
                else:
                    continue
            
            controller = ctask.getAssignedRobotController()
            if ii == start:
                start_robot_id = controller.getControlledRobotID()
            task = ctask.getNextUncompletedTask()
            self._logger.debug('Time={0:.2f} TCS try to schedule the task ({1},{2}) assigned to {3}'.\
                              format(self._timer.getTime(), task.getSrc(), task.getDst(), controller.getControlledRobotID())) 
            last_waypoint = controller.getLastWaypoint()
            next_waypoint = controller.getNextWaypoint()
            if last_waypoint != task.getDst():
                if controller.isOnEdge():
                    if next_waypoint == task.getDst():
                        paths = [[last_waypoint, next_waypoint]]
                    else:
                        paths = []
                        temp_paths = self._router.getKShortestPaths(next_waypoint, task.getDst())
                        for i in range(len(temp_paths)):
                            if last_waypoint not in temp_paths[i]: # we don't want paths going back to last_waypoint
                                paths.append([last_waypoint] + temp_paths[i])
                else:
                    paths = self._router.getKShortestPaths(last_waypoint, task.getDst())
                path, waypoint_time = self._scheduler.schedule(task, controller, paths, robot2index, verbose=verbose)
                if path == None or waypoint_time == None:
                    self._scheduler.restore()
                    #return {start_robot_id:(None, None)} # 如果任何一个任务调度失败，则不做任何实质性地调度。仅返回下标为start的任务调度失败的信息。
                    return {}
                schedule_results[controller.getControlledRobotID()] = (path, waypoint_time)
        return schedule_results
        
    def doActualSchedule(self, schedule_results):
        """
        进行实际调度。即，通过发送消息，将调度结果最终发送给机器人。
        :param schedule_results: 调度结果。用一个字典表示，比如{0:(['1','2','3'], [0.1, 1.0, 1.9]), 1:[None, None]}
                                 表示机器人0调度的路径是['1','2','3']，到达这三个点的时间依次为[0.1, 1.0, 1.9]。机器人1调度失败。
        """
        for robot_id, schedule_result in schedule_results.items():
            path, waypoint_time = schedule_result
            controller = self._robots_manager.getRobotController(robot_id)
            ctask = controller.getCompositeTask()
            task = ctask.getNextUncompletedTask()
            
            if path is not None and waypoint_time is not None:
                self._logger.info('Time={0:.2f} TCS scheduling robot {1}:\n  path {2}\n  time {3}.'.\
                                  format(self._timer.getTime(), controller.getControlledRobotID(), path, waypoint_time))
                controller.schedule(path, waypoint_time)
                ctask.setFirstTaskProcessing()
                task.setReleaseTime(waypoint_time[-1])
            else:
                self._logger.info('Time={0:.2f} TCS scheduling robot {1} failed.'.\
                              format(self._timer.getTime(), robot_id))
                controller.pause()
                ctask.setFirstTaskPending()
                task.setReleaseTime(float('inf'))
                self.updatePriority()
                        
    def _indexOfNewTask(self, start=0):
        sorted_ctasks = list(self._ctask_queue)
        for ii in range(start, len(sorted_ctasks)):
            ctask = sorted_ctasks[ii]
            task = ctask.getNextUncompletedTask()
            if ctask.isAssigned() and task.isNotScheduled():
                return ii
        return -1
        
    def _indexOfPendingTask(self, start=0):
        ctasks = list(self._ctask_queue)
        for ii in range(start, len(ctasks)):
            ctask = ctasks[ii]
            task = ctask.getNextUncompletedTask()
            if task.isPending():
                return ii
        return -1
        
    def _indexOfProcessingTask(self, start=0):
        ctasks = list(self._ctask_queue)
        for ii in range(start, len(ctasks)):
            ctask = ctasks[ii]
            task = ctask.getNextUncompletedTask()
            if task.isProcessing():
                return ii
        return -1

    def _indexOfProcessingOrPendingTask(self, start=0, skip=None, reverse=False):
        ctasks = list(self._ctask_queue)
        if start < 0:
            start += len(ctasks)
        if skip is not None:
            if skip < 0:
                skip += len(ctasks)
                
        if reverse:
            rg = range(start, -1, -1)
        else:
            rg = range(start, len(ctasks))
            
        for ii in rg:
            if skip is not None and skip == ii:
                continue
                
            ctask = ctasks[ii]
            task = ctask.getNextUncompletedTask()
            if task.isPending() or task.isProcessing():
                return ii
        return -1
    
    def _updateRobot2Index(self):
        ctasks = list(self._ctask_queue)
        self._robot2index = {}
        i = 0
        for ctask in self._ctask_queue:
            if ctask.isAssigned():
                self._robot2index[ctask.getAssignedRobotController().getControlledRobotID()] = i
                i += 1
                
    def reschedule(self, event, verbose=False, planning=False):
        """
        对优先级小于等于priority的所有任务进行重新调度。
        """        
        sorted_ctasks = list(self._ctask_queue)
        self._updateRobot2Index()
        
        for ctask in sorted_ctasks:
            if ctask.isAssigned():
                robot_id = ctask.getAssignedRobotController().getControlledRobotID()
                task = ctask.getNextUncompletedTask()
                self._logger.debug('robot {0} task state {1}'.format(robot_id, task.getState()))
        
        if event == TransportationControlSystem.EVENT_NEW_TASK_CREATED:
            # find the new task
            start = 0
            schedule_results = {}
            while True:
                ctask_idx = self._indexOfNewTask(start)
                self._logger.debug('Reschedule because of new task {0} arrived.'.format(ctask_idx))
                if ctask_idx >= 0:
                    results = self._rescheduleTasks(self._robot2index, start=ctask_idx, reschedule_lower_priority_tasks=True, verbose=verbose)
                    start = ctask_idx + 1
                    if results:
                        schedule_results.update(results)
                else:
                    break
            if schedule_results and (not planning):
                self.doActualSchedule(schedule_results)
        elif event == TransportationControlSystem.EVENT_ROBOT_POS_UPDATED:
            # find the first pending task
            start = 0
            schedule_results = {}
            while start < len(sorted_ctasks):
                ctask_idx = self._indexOfPendingTask(start)
                self._logger.debug('Reschedule ctask {0} because robot position updated.'.format(ctask_idx))
                if ctask_idx >= 0:
                    results = self._rescheduleTasks(self._robot2index, start=ctask_idx, reschedule_lower_priority_tasks=True, verbose=verbose)
                    start = ctask_idx + 1
                    if results:
                        schedule_results.update(results)
                else:
                    break
            if schedule_results and (not planning):
                self.doActualSchedule(schedule_results)
                    
    def pendingCompositeTasks(self, start=0):
        sorted_ctasks = list(self._ctask_queue)
        pending_ctasks = []
        for ii in range(start, len(sorted_ctasks)):
            ctask = sorted_ctasks[ii]
            task = ctask.getNextUncompletedTask()
            if task.isPending():
                pending_ctasks.append(ii)
        return pending_ctasks
                    
    def avgResponseTime(self, task_creator='all'):
        """
        当前时刻，所有未完成的任务的平均等待时间（从发起时间计起，到时间窗中所指示的结束时间为止）。
        :param task_creator: 只对该创建者创建的任务统计平均等待时间。可设置的值为'all', 'user'或者'system'.
        :return: avg_response_time, num_pending_tasks，正在进行中的任务的平均等待时间以及
        """
        time_sum = 0.0
        num = 0
        num_pending_tasks = 0
        for ctask in self._ctask_queue:
            #print('Trying to update ', ctask.getAssignedRobotController().getControlledRobotID(), '\'s task priority')
            if (not ctask.isEmpty()):
                task = ctask.getNextUncompletedTask()
                if task.isPending() or task.isProcessing():
                    if task_creator == 'all' or \
                      ((task_creator == 'user' or task_creator == 'system') and task.getCreator() == task_creator):
                        release_time = task.getReleaseTime()
                        if release_time == float('inf'): # 如果任何一个任务的预期结束时间为无穷，则平均等待时间为无穷
                            num_pending_tasks += 1
                        else:
                            time_sum += release_time - ctask.getCreateTime()
                            num += 1
        if num == 0:
            return 0.0, num_pending_tasks
        else:
            return time_sum/(1.0*num), num_pending_tasks